#include <ompl/contrib/hplanner/approx_cspace_oracle.h>

namespace ompl
{
namespace base
{


#if CONSIDER_NONEUCLIDEAN
  std::pair<float, float> ApproxCSpaceOracleBase::computeEmbedParamFromDataset() const
  {
    std::pair<float, float> embed_c;

    bool* types = svt_.getTypes();

    float param = 0;
    unsigned int nNegativeRawState = getNumNegativeRawState();

    for(unsigned int i = 0; i < nNegativeRawState; ++i)
    {
      float tmp = 0;
      float* state = getNegativeRawState(i);
      for(unsigned int j = 0; j < dim_; ++j)
      {
        float value = state[j];
        if(types[j])
          tmp += value * value;
        else
        {
          if(value < 0)
            tmp += (2 * ApproxCSpaceOracle_PI - value) * (2 * ApproxCSpaceOracle_PI - value);
          else
            tmp += (value - 2 * ApproxCSpaceOracle_PI) * (value - 2 * ApproxCSpaceOracle_PI);
        }
      }
      if(tmp > param) param = tmp;
    }
    embed_c.first = param;

    if(!getUsePositiveState())
    {
      return embed_c;
    }

    param = 0;
    unsigned int nPositiveRawState = getNumPositiveRawState();
    for(unsigned int i = 0; i < nPositiveRawState; ++i)
    {
      float tmp = 0;
      float* state = getPositiveRawState(i);
      for(unsigned int j = 0; j < dim_; ++j)
      {
        float value = state[j];
        if(types[j])
          tmp += value * value;
        else
        {
          if(value < 0)
            tmp += (2 * ApproxCSpaceOracle_PI - value) * (2 * ApproxCSpaceOracle_PI - value);
          else
            tmp += (value - 2 * ApproxCSpaceOracle_PI) * (value - 2 * ApproxCSpaceOracle_PI);
        }
      }
      if(tmp > param) param = tmp;

    }
    embed_c.second = param;

    return embed_c;
  }
#else
  std::pair<float, float> ApproxCSpaceOracleBase::computeEmbedParamFromDataset() const
  {
    cspace_learning::dotprod<float> l2dot(dim_);

    std::pair<float, float> embed_c;

    float param = 0;
    unsigned int nNegativeRawState = getNumNegativeRawState();

    for(unsigned int i = 0; i < nNegativeRawState; ++i)
    {
      float* state = getNegativeRawState(i);
      float tmp = l2dot(state, state);
      if(tmp > param) param = tmp;
    }
    embed_c.first = param;

    if(!getUsePositiveState())
    {
      return embed_c;
    }

    param = 0;
    unsigned int nPositiveRawState = getNumPositiveRawState();

    for(unsigned int i = 0; i < nPositiveRawState; ++i)
    {
      float* state = getPositiveRawState(i);
      float tmp = l2dot(state, state);
      if(tmp > param) param = tmp;
    }
    embed_c.second = param;

    return embed_c;
  }
#endif

#if CONSIDER_NONEUCLIDEAN
  std::pair<float, float> ApproxCSpaceOracleBase::computeEmbedParamDefault() const
  {
    bool* types = svt_.getTypes();
    float tmp = 0;
    float* upper_bound = svt_.getUpperBound();

    for(unsigned int i = 0; i < dim_; ++i)
    {
      if(types[i])
        tmp += upper_bound[i] * upper_bound[i];
      else
        tmp += 4 * upper_bound[i] * upper_bound[i];
    }

    float param = tmp;

    tmp = 0;
    float* lower_bound = svt_.getLowerBound();
    for(unsigned int i = 0; i < dim_; ++i)
    {
      if(types[i])
        tmp += lower_bound[i] * lower_bound[i];
      else
        tmp += 4 * lower_bound[i] * lower_bound[i];
    }

    if(tmp > param) param = tmp;

    return std::make_pair(param, param);
  }
#else
  void ApproxCSpaceOracleBase::computeEmbedParamDefault() const
  {
    cspace_learning::dotprod<float> l2dot(dim_);

    float* upper_bound = svt_.getUpperBound();
    float* lower_bound = svt_.getLowerBound();

    float param = l2dot(upper_bound, upper_bound);
    float param2 = l2dot(lower_bound, lower_bound);

    if(param2 > param) param = param2;

    return std::make_pair(param, param);
  }
#endif







ApproxCSpaceOracleBasic::ApproxCSpaceOracleBasic(const SpaceInformationPtr& si,
                                                 unsigned int K1, float R1, unsigned int M1, unsigned int L1, unsigned int H1, float W1,
                                                 unsigned int K2, float R2, unsigned int M2, unsigned int L2, unsigned int H2, float W2,
                                                 bool usePositiveState)
                                                  : ApproxCSpaceOracleBase(si, usePositiveState)
{
  embed_dim_ = (dim_ + 2) * (dim_ + 3) / 2;

  line_l2_ = new L2sqr(embed_dim_);
  point_l2_ = new L2sqr(dim_);
  l2dot_ = new Dotprod(dim_);
  embed_ = new Embedding(dim_);
  line_eval_ = new PointToLineEval(dim_);
  point_eval_ = new PointEval(*point_l2_);


  collision_dataset_ = new Dataset(dim_);
  collision_accessor_ = new Accessor(*collision_dataset_);
  collision_line_scanner_ = new LineScanner(*collision_accessor_, line_eval_, K1, R1);
  collision_point_scanner_ = new PointScanner(*collision_accessor_, point_eval_, K1, R1);

  Index::Parameter collision_line_param;
  collision_line_param.repeat = M1;
  collision_line_param.range = H1;
  collision_line_param.W = W1;
  collision_line_param.dim = embed_dim_;
  collision_line_index_.init(collision_line_param, rng_, L1);

  Index::Parameter collision_point_param;
  collision_point_param.repeat = M1;
  collision_point_param.range = H1;
  collision_point_param.W = W1;
  collision_point_param.dim = dim_;
  collision_point_index_.init(collision_point_param, rng_, L1);


  if(getUsePositiveState())
  {
    collision_free_dataset_ = new Dataset(dim_);
    collision_free_accessor_ = new Accessor(*collision_free_dataset_);
    collision_free_line_scanner_ = new LineScanner(*collision_free_accessor_, line_eval_, K2, R2);
    collision_free_point_scanner_ = new PointScanner(*collision_free_accessor_, point_eval_, K2, R2);

    Index::Parameter collision_free_line_param;
    collision_free_line_param.repeat = M2;
    collision_free_line_param.range = H2;
    collision_free_line_param.W = W2;
    collision_free_line_param.dim = embed_dim_;
    collision_free_line_index_.init(collision_free_line_param, rng_, L2);

    Index::Parameter collision_free_point_param;
    collision_free_point_param.repeat = M2;
    collision_free_point_param.range = H2;
    collision_free_point_param.W = W2;
    collision_free_point_param.dim = dim_;
    collision_free_point_index_.init(collision_free_point_param, rng_, L2);
  }
  else
  {
    collision_free_dataset_ = NULL;
    collision_free_accessor_ = NULL;
    collision_free_line_scanner_ = NULL;
    collision_free_point_scanner_ = NULL;
  }


  line_direction_cache_ = new float[dim_];
  line_origin_cache_ = new float[dim_];
  transform_state_cache_ = new float[dim_];
  transformed_state_1_cache_ = new float[dim_];
  transformed_state_2_cache_ = new float[dim_];

  std::pair<float, float> embed_c = computeEmbedParamDefault();
  collision_embed_c_ = embed_c.first;
  collision_free_embed_c_ = embed_c.second;
}

ApproxCSpaceOracleBasic::~ApproxCSpaceOracleBasic()
{
  delete collision_accessor_;
  delete collision_dataset_;
  delete collision_line_scanner_;
  delete collision_point_scanner_;

  if(getUsePositiveState())
  {
    delete collision_free_accessor_;
    delete collision_free_dataset_;
    delete collision_free_line_scanner_;
    delete collision_free_point_scanner_;
  }

  delete line_l2_;
  delete point_l2_;
  delete l2dot_;
  delete embed_;

  delete line_eval_;
  delete point_eval_;

  delete [] line_direction_cache_;
  delete [] line_origin_cache_;
  delete [] transform_state_cache_;
  delete [] transformed_state_1_cache_;
  delete [] transformed_state_2_cache_;

  for(unsigned int i = 0; i < collision_samples_.size(); ++i)
    si_->freeState(collision_samples_[i]);
  for(unsigned int i = 0; i < collision_free_samples_.size(); ++i)
  {
    si_->freeState(collision_free_samples_[i]);
  }
}


void ApproxCSpaceOracleBasic::insert(const State* state, bool collision)
{
  if(!collision && !getUsePositiveState())
    return;

  float* transformed_state = svt_.transform(state);

  if(collision)
  {
    collision_dataset_->insert(transformed_state, dim_, 1);
    collision_samples_.push_back(si_->cloneState(state));
  }
  else
  {
    collision_free_dataset_->insert(transformed_state, dim_, 0);
    collision_free_samples_.push_back(si_->cloneState(state));
  }
}


#if CONSIDER_NONEUCLIDEAN

  void ApproxCSpaceOracleBasic::insertAndIndex(const State* state, bool collision)
  {
    if(!collision && !getUsePositiveState())
      return;

    insert(state, collision);

    if(collision)
      index(collision_dataset_->size() - 1, collision);
    else
      index(collision_free_dataset_->size() - 1, collision);
  }

  void ApproxCSpaceOracleBasic::index(unsigned int id, bool collision)
  {
    float* state = NULL;

    if(collision)
    {
      state =  (*collision_dataset_)[id];
      embed_->t() = sqrt(collision_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_point_index_.insert(id, state);
    }
    else
    {
      state = (*collision_free_dataset_)[id];
      embed_->t() = sqrt(collision_free_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_free_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_free_point_index_.insert(id, state);
    }

    unsigned int noneu_bits = svt_.nonEuclideanDim();
    bool* types = svt_.getTypes();

    unsigned long N = 0;

    boost::dynamic_bitset<> bits(noneu_bits, N);
    do
    {
      N++;
      bits = boost::dynamic_bitset<>(noneu_bits, N);

      unsigned noneu_bit_id = 0;
      for(unsigned int i = 0; i < dim_; ++i)
      {
        if(types[i])
          transform_state_cache_[i] = state[i];
        else
        {
          if(bits[noneu_bit_id] == 1)
          {
            if(state[i] < 0) transform_state_cache_[i] = state[i] +  2 * ApproxCSpaceOracle_PI;
            else transform_state_cache_[i] = state[i] - 2 * ApproxCSpaceOracle_PI;
          }
          else
            transform_state_cache_[i] = state[i];
          noneu_bit_id++;
        }
      }


      if(collision)
      {
        embed_->t() = sqrt(collision_embed_c_ - 1 - (*l2dot_)(transform_state_cache_, transform_state_cache_));
        collision_line_index_.insert(id, embed_->embedded_datapoint(transform_state_cache_));
        collision_point_index_.insert(id, transform_state_cache_);
      }
      else
      {
        embed_->t() = sqrt(collision_free_embed_c_ - 1 - (*l2dot_)(transform_state_cache_, transform_state_cache_));
        collision_free_line_index_.insert(id, embed_->embedded_datapoint(transform_state_cache_));
        collision_free_point_index_.insert(id, transform_state_cache_);
      }
    }
    while(bits.count() != noneu_bits);
  }
#else

  void ApproxCSpaceOracleBasic::insertAndIndex(const State* state, bool collision)
  {
    if(!collision && getUsePositiveState())
      return;

    insert(state, collision);

    if(collision)
      index(collision_dataset_->size() - 1, collision);
    else
      index(collision_free_dataset_->size() - 1, collision);
  }

  void ApproxCSpaceOracleBasic::index(unsigned int id, bool collision)
  {
    if(collision)
    {
      float* state = (*collision_dataset_)[id];
      embed_->t() = sqrt(collision_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_point_index_.insert(id, state);
    }
    else
    {
      float* state = (*collision_free_dataset_)[id];
      embed_->t() = sqrt(collision_free_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_free_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_free_point_index_.insert(id, state);
    }
  }

#endif

  void ApproxCSpaceOracleBasic::pointKNNQuery(const State* state, bool collision,
                                              std::vector<std::pair<unsigned int, float> >& res) const
  {
    if(!collision && !getUsePositiveState())
      return;

    float* transformed_state = svt_.transform(state);

    if(collision)
    {
      collision_accessor_->reset();
      collision_point_scanner_->reset(transformed_state);
      collision_point_index_.query(*collision_point_scanner_);

      unsigned int nTopk = collision_point_scanner_->topk().nTopk();
      res.resize(nTopk);
      for(unsigned int i = 0; i < nTopk; ++i)
      {
        res[i].first = collision_point_scanner_->topk()[i].key;
        res[i].second = collision_point_scanner_->topk()[i].dist;
      }
    }
    else
    {
      collision_free_accessor_->reset();
      collision_free_point_scanner_->reset(transformed_state);
      collision_free_point_index_.query(*collision_free_point_scanner_);

      unsigned int nTopk = collision_free_point_scanner_->topk().nTopk();
      res.resize(nTopk);
      for(unsigned int i = 0; i < nTopk; ++i)
      {
        res[i].first = collision_free_point_scanner_->topk()[i].key;
        res[i].second = collision_free_point_scanner_->topk()[i].dist;
      }
    }
  }

  void ApproxCSpaceOracleBasic::lineKNNQuery(const State* s1, const State* s2, bool collision,
                                             std::vector<std::pair<unsigned int, float> >& res,
                                             float* line_origin, float* line_direction, float* line_len) const
  {
    if(!collision && !getUsePositiveState())
      return;

    memcpy(transformed_state_1_cache_, svt_.transform(s1), sizeof(float) * dim_);
    memcpy(transformed_state_2_cache_, svt_.transform(s2), sizeof(float) * dim_);
    memcpy(line_origin_cache_, transformed_state_1_cache_, sizeof(float) * dim_);

    float line_len_cache_ = 0;
    for(unsigned int i = 0; i < dim_; ++i)
    {
      line_direction_cache_[i] = transformed_state_2_cache_[i] - transformed_state_1_cache_[i];
      line_len_cache_ += line_direction_cache_[i] * line_direction_cache_[i];
    }

    line_len_cache_ = sqrt(line_len_cache_);

    for(unsigned int i = 0; i < dim_; ++i)
      line_direction_cache_[i] /= line_len_cache_;

    if(collision)
    {
      collision_accessor_->reset();
      line_eval_->reset(line_direction_cache_, line_origin_cache_);
      collision_line_scanner_->reset(embed_->embedded_line(line_direction_cache_, line_origin_cache_));
      collision_line_index_.query(*collision_line_scanner_);

      cspace_learning::Topk<typename LineScanner::Key> tmp;
      float* item = new float[dim_];

      unsigned int nTopk = collision_line_scanner_->topk().nTopk();

      for(unsigned int i = 0; i < nTopk; ++i)
      {
        memcpy(item, (*collision_dataset_)[collision_line_scanner_->topk()[i].key], sizeof(float) * dim_);
        for(unsigned int j = 0; j < dim_; ++j)
          item[j] -= line_origin_cache_[j];
        float proj = (*l2dot_)(item, line_direction_cache_);
        if(proj > 0 && proj < line_len_cache_)
        {
          res.push_back(std::make_pair(collision_line_scanner_->topk()[i].key, collision_line_scanner_->topk()[i].dist));
        }
      }
      delete [] item;
    }
    else
    {
      collision_free_accessor_->reset();
      line_eval_->reset(line_direction_cache_, line_origin_cache_);
      collision_free_line_scanner_->reset(embed_->embedded_line(line_direction_cache_, line_origin_cache_));
      collision_free_line_index_.query(*collision_free_line_scanner_);

      cspace_learning::Topk<typename LineScanner::Key> tmp;
      float* item = new float[dim_];

      unsigned int nTopk = collision_free_line_scanner_->topk().nTopk();

      for(unsigned int i = 0; i < nTopk; ++i)
      {
        memcpy(item, (*collision_free_dataset_)[collision_free_line_scanner_->topk()[i].key], sizeof(float) * dim_);
        for(unsigned int j = 0; j < dim_; ++j)
          item[j] -= line_origin_cache_[j];
        float proj = (*l2dot_)(item, line_direction_cache_);
        if(proj > 0 && proj < line_len_cache_)
        {
          res.push_back(std::make_pair(collision_free_line_scanner_->topk()[i].key, collision_free_line_scanner_->topk()[i].dist));
        }
      }
      delete [] item;
    }

    if(line_origin) memcpy(line_origin, line_origin_cache_, sizeof(float) * dim_);
    if(line_direction) memcpy(line_direction, line_direction_cache_, sizeof(float) * dim_);
    if(line_len) *line_len = line_len_cache_;
  }



  ApproxCSpaceOracleMultiprobe::ApproxCSpaceOracleMultiprobe(const SpaceInformationPtr& si,
                                                             unsigned int K1, float R1, unsigned int M1, unsigned int L1, unsigned int H1, float W1,
                                                             unsigned int K2, float R2, unsigned int M2, unsigned int L2, unsigned int H2, float W2,
                                                             bool usePositiveState) : ApproxCSpaceOracleBase(si, usePositiveState)
  {
    multiprobe_num_ = 5; // multiprobe

    embed_dim_ = (dim_ + 2) * (dim_ + 3) / 2;

    line_l2_ = new L2sqr(embed_dim_);
    point_l2_ = new L2sqr(dim_);
    l2dot_ = new Dotprod(dim_);
    embed_ = new Embedding(dim_);
    line_eval_ = new PointToLineEval(dim_);
    point_eval_ = new PointEval(*point_l2_);

    collision_dataset_ = new Dataset(dim_);
    collision_accessor_ = new Accessor(*collision_dataset_);
    collision_line_scanner_ = new LineScanner(*collision_accessor_, line_eval_, K1, R1);
    collision_point_scanner_ = new PointScanner(*collision_accessor_, point_eval_, K1, R1);

    Index::Parameter collision_line_param;
    collision_line_param.repeat = M1;
    collision_line_param.range = H1;
    collision_line_param.W = W1;
    collision_line_param.dim = embed_dim_;
    collision_line_index_.init(collision_line_param, rng_, L1);

    Index::Parameter collision_point_param;
    collision_point_param.repeat = M1;
    collision_point_param.range = H1;
    collision_point_param.W = W1;
    collision_point_param.dim = dim_;
    collision_point_index_.init(collision_point_param, rng_, L1);

    if(getUsePositiveState())
    {
      collision_free_dataset_ = new Dataset(dim_);
      collision_free_accessor_ = new Accessor(*collision_free_dataset_);
      collision_free_line_scanner_ = new LineScanner(*collision_free_accessor_, line_eval_, K2, R2);
      collision_free_point_scanner_ = new PointScanner(*collision_free_accessor_, point_eval_, K2, R2);

      Index::Parameter collision_free_line_param;
      collision_free_line_param.repeat = M2;
      collision_free_line_param.range = H2;
      collision_free_line_param.W = W2;
      collision_free_line_param.dim = embed_dim_;
      collision_free_line_index_.init(collision_free_line_param, rng_, L2);

      Index::Parameter collision_free_point_param;
      collision_free_point_param.repeat = M2;
      collision_free_point_param.range = H2;
      collision_free_point_param.W = W2;
      collision_free_point_param.dim = dim_;
      collision_free_point_index_.init(collision_free_point_param, rng_, L2);
    }
    else
    {
      collision_free_dataset_ = NULL;
      collision_free_accessor_ = NULL;
      collision_free_line_scanner_ = NULL;
      collision_free_point_scanner_ = NULL;
    }


    line_direction_cache_ = new float[dim_];
    line_origin_cache_ = new float[dim_];
    transform_state_cache_ = new float[dim_];
    transformed_state_1_cache_ = new float[dim_];
    transformed_state_2_cache_ = new float[dim_];

    std::pair<float, float> embed_c = computeEmbedParamDefault();
    collision_embed_c_ = embed_c.first;
    collision_free_embed_c_ = embed_c.second;
  }


  ApproxCSpaceOracleMultiprobe::~ApproxCSpaceOracleMultiprobe()
  {
    delete collision_accessor_;
    delete collision_dataset_;
    delete collision_line_scanner_;
    delete collision_point_scanner_;

    if(getUsePositiveState())
    {
      delete collision_free_accessor_;
      delete collision_free_dataset_;
      delete collision_free_line_scanner_;
      delete collision_free_point_scanner_;
    }

    delete line_l2_;
    delete point_l2_;
    delete l2dot_;
    delete embed_;

    delete line_eval_;
    delete point_eval_;

    delete [] line_direction_cache_;
    delete [] line_origin_cache_;
    delete [] transform_state_cache_;
    delete [] transformed_state_1_cache_;
    delete [] transformed_state_2_cache_;

    for(unsigned int i = 0; i < collision_samples_.size(); ++i)
      si_->freeState(collision_samples_[i]);
    for(unsigned int i = 0; i < collision_free_samples_.size(); ++i)
    {
      si_->freeState(collision_free_samples_[i]);
    }
  }

  void ApproxCSpaceOracleMultiprobe::insert(const State* state, bool collision)
  {
    if(!collision && !getUsePositiveState())
      return;

    float* transformed_state = svt_.transform(state);

    if(collision)
    {
      collision_dataset_->insert(transformed_state, dim_, 1);
      collision_samples_.push_back(si_->cloneState(state));
    }
    else
    {
      collision_free_dataset_->insert(transformed_state, dim_, 0);
      collision_free_samples_.push_back(si_->cloneState(state));
    }
  }


#if CONSIDER_NONEUCLIDEAN

  void ApproxCSpaceOracleMultiprobe::insertAndIndex(const State* state, bool collision)
  {
    if(!collision && !getUsePositiveState())
      return;

    insert(state, collision);

    if(collision)
      index(collision_dataset_->size() - 1, collision);
    else
      index(collision_free_dataset_->size() - 1, collision);
  }


  void ApproxCSpaceOracleMultiprobe::index(unsigned int id, bool collision)
  {
    float* state = NULL;

    if(collision)
    {
      state =  (*collision_dataset_)[id];
      embed_->t() = sqrt(collision_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_point_index_.insert(id, state);
    }
    else
    {
      state = (*collision_free_dataset_)[id];
      embed_->t() = sqrt(collision_free_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_free_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_free_point_index_.insert(id, state);
    }

    unsigned int noneu_bits = svt_.nonEuclideanDim();
    bool* types = svt_.getTypes();

    unsigned long N = 0;

    boost::dynamic_bitset<> bits(noneu_bits, N);
    do
    {
      N++;
      bits = boost::dynamic_bitset<>(noneu_bits, N);

      unsigned noneu_bit_id = 0;
      for(unsigned int i = 0; i < dim_; ++i)
      {
        if(types[i])
          transform_state_cache_[i] = state[i];
        else
        {
          if(bits[noneu_bit_id] == 1)
          {
            if(state[i] < 0) transform_state_cache_[i] = state[i] +  2 * ApproxCSpaceOracle_PI;
            else transform_state_cache_[i] = state[i] - 2 * ApproxCSpaceOracle_PI;
          }
          else
            transform_state_cache_[i] = state[i];
          noneu_bit_id++;
        }
      }

      if(collision)
      {
        embed_->t() = sqrt(collision_embed_c_ - 1 - (*l2dot_)(transform_state_cache_, transform_state_cache_));
        collision_line_index_.insert(id, embed_->embedded_datapoint(transform_state_cache_));
        collision_point_index_.insert(id, transform_state_cache_);
      }
      else
      {
        embed_->t() = sqrt(collision_free_embed_c_ - 1 - (*l2dot_)(transform_state_cache_, transform_state_cache_));
        collision_free_line_index_.insert(id, embed_->embedded_datapoint(transform_state_cache_));
        collision_free_point_index_.insert(id, transform_state_cache_);
      }
    }
    while(bits.count() != noneu_bits);
  }
#else
  void ApproxCSpaceOracleMultiprobe::insertAndIndex(const State* state, bool collision)
  {
    if(!collision && !getUsePositiveState())
      return;

    insert(state, collision);

    if(collision)
      index(collision_dataset_->size() - 1, collision);
    else
      index(collision_free_dataset_->size() - 1, collision);
  }

  void ApproxCSpaceOracleMultiprobe::index(unsigned int id, bool collision)
  {
    if(collision)
    {
      float* state = (*collision_dataset_)[id];
      embed_->t() = sqrt(collision_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_point_index_.insert(id, state);
    }
    else
    {
      float* state = (*collision_free_dataset_)[id];
      embed_->t() = sqrt(collision_free_embed_c_ - 1 - (*l2dot_)(state, state));
      collision_free_line_index_.insert(id, embed_->embedded_datapoint(state));
      collision_free_point_index_.insert(id, state);
    }
  }
#endif


  void ApproxCSpaceOracleMultiprobe::pointKNNQuery(const State* state, bool collision,
                                                   std::vector<std::pair<unsigned int, float> >& res) const
  {
    if(!collision && !getUsePositiveState())
      return;

    float* transformed_state = svt_.transform(state);

    if(collision)
    {
      collision_accessor_->reset();
      collision_point_scanner_->reset(transformed_state);
      collision_point_index_.query(multiprobe_num_, *collision_point_scanner_);

      unsigned int nTopk = collision_point_scanner_->topk().nTopk();
      res.resize(nTopk);
      for(unsigned int i = 0; i < nTopk; ++i)
      {
        res[i].first = collision_point_scanner_->topk()[i].key;
        res[i].second = collision_point_scanner_->topk()[i].dist;
      }
    }
    else
    {
      collision_free_accessor_->reset();
      collision_free_point_scanner_->reset(transformed_state);
      collision_free_point_index_.query(multiprobe_num_, *collision_free_point_scanner_);

      unsigned int nTopk = collision_free_point_scanner_->topk().nTopk();
      res.resize(nTopk);
      for(unsigned int i = 0; i < nTopk; ++i)
      {
        res[i].first = collision_free_point_scanner_->topk()[i].key;
        res[i].second = collision_free_point_scanner_->topk()[i].dist;
      }
    }
  }

  void ApproxCSpaceOracleMultiprobe::lineKNNQuery(const State* s1, const State* s2, bool collision,
                                                  std::vector<std::pair<unsigned int, float> >& res,
                                                  float* line_origin, float* line_direction, float* line_len) const
  {
    if(!collision && !getUsePositiveState())
      return;

    memcpy(transformed_state_1_cache_, svt_.transform(s1), sizeof(float) * dim_);
    memcpy(transformed_state_2_cache_, svt_.transform(s2), sizeof(float) * dim_);
    memcpy(line_origin_cache_, transformed_state_1_cache_, sizeof(float) * dim_);

    float line_len_cache_ = 0;
    for(unsigned int i = 0; i < dim_; ++i)
    {
      line_direction_cache_[i] = transformed_state_2_cache_[i] - transformed_state_1_cache_[i];
      line_len_cache_ += line_direction_cache_[i] * line_direction_cache_[i];
    }

    line_len_cache_ = sqrt(line_len_cache_);


    for(unsigned int i = 0; i < dim_; ++i)
      line_direction_cache_[i] /= line_len_cache_;

    if(collision)
    {
      collision_accessor_->reset();
      line_eval_->reset(line_direction_cache_, line_origin_cache_);
      collision_line_scanner_->reset(embed_->embedded_line(line_direction_cache_, line_origin_cache_));
      collision_line_index_.query(multiprobe_num_, *collision_line_scanner_);

      cspace_learning::Topk<typename LineScanner::Key> tmp;
      float* item = new float[dim_];

      unsigned int nTopk = collision_line_scanner_->topk().nTopk();

      for(unsigned int i = 0; i < nTopk; ++i)
      {
        memcpy(item, (*collision_dataset_)[collision_line_scanner_->topk()[i].key], sizeof(float) * dim_);
        for(unsigned int j = 0; j < dim_; ++j)
          item[j] -= line_origin_cache_[j];
        float proj = (*l2dot_)(item, line_direction_cache_);
        if(proj > 0 && proj < line_len_cache_)
        {
          res.push_back(std::make_pair(collision_line_scanner_->topk()[i].key, collision_line_scanner_->topk()[i].dist));
        }
      }
      delete [] item;
    }
    else
    {
      collision_free_accessor_->reset();
      line_eval_->reset(line_direction_cache_, line_origin_cache_);
      collision_free_line_scanner_->reset(embed_->embedded_line(line_direction_cache_, line_origin_cache_));
      collision_free_line_index_.query(multiprobe_num_, *collision_free_line_scanner_);

      cspace_learning::Topk<typename LineScanner::Key> tmp;
      float* item = new float[dim_];

      unsigned int nTopk = collision_free_line_scanner_->topk().nTopk();

      for(unsigned int i = 0; i < nTopk; ++i)
      {
        memcpy(item, (*collision_free_dataset_)[collision_free_line_scanner_->topk()[i].key], sizeof(float) * dim_);
        for(unsigned int j = 0; j < dim_; ++j)
          item[j] -= line_origin_cache_[j];
        float proj = (*l2dot_)(item, line_direction_cache_);
        if(proj > 0 && proj < line_len_cache_)
        {
          res.push_back(std::make_pair(collision_free_line_scanner_->topk()[i].key, collision_free_line_scanner_->topk()[i].dist));
        }
      }
      delete [] item;
    }

    if(line_origin) memcpy(line_origin, line_origin_cache_, sizeof(float) * dim_);
    if(line_direction) memcpy(line_direction, line_direction_cache_, sizeof(float) * dim_);
    if(line_len) *line_len = line_len_cache_;
  }

}
}
